from keras import Sequential
import os
import tensorflow as tf
from keras.layers import Dense

if __name__ == "__main__":
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

    # gpus = tf.config.experimental.list_physical_devices('GPU')
    # if gpus:
    #     try:
    #         # 设置内存自增长
    #         for gpu in gpus:
    #             tf.config.experimental.set_memory_growth(gpu, True)
    #         print('-------------已设置完GPU内存自增长--------------')
    #
    #         # 指定使用第二块GPU
    #         # tf.config.experimental.set_visible_devices(gpus[1], 'GPU')
    #
    #         # 获取逻辑GPU个数
    #         logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    #         print('逻辑GPU个数为：', len(logical_gpus))
    #     except RuntimeError as e:
    #         # 打印异常信息
    #         print(e)
    model = Sequential(
        [
            Dense(256, activation='relu'),
            Dense(256, activation='relu'),
            Dense(256, activation='relu'),
            Dense(10, activation='relu'),
        ]
    )

    model.build(input_shape=[4, 784])

    model.summary()